#!/usr/bin/env python3
from flask import Flask, render_template, request, jsonify
import json
import subprocess
import threading
import time
from datetime import datetime

app = Flask(__name__)

task_status = {
    'running': False,
    'last_run': None,
    'results': []
}

def load_config():
    try:
        with open('servers.json', 'r', encoding='utf-8') as f:
            return json.load(f)
    except:
        return {'servers': []}

def save_config(config):
    with open('servers.json', 'w', encoding='utf-8') as f:
        json.dump(config, f, ensure_ascii=False, indent=2)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/servers', methods=['GET'])
def get_servers():
    config = load_config()
    servers_safe = []
    for server in config['servers']:
        servers_safe.append({
            'host': server['host'],
            'username': server['username'],
            'port': server.get('port', 22)
        })
    return jsonify(servers_safe)

@app.route('/api/servers', methods=['POST'])
def add_server():
    data = request.json
    config = load_config()
    config['servers'].append(data)
    save_config(config)
    return jsonify({'success': True})

@app.route('/api/servers/<host>', methods=['DELETE'])
def delete_server(host):
    config = load_config()
    config['servers'] = [s for s in config['servers'] if s['host'] != host]
    save_config(config)
    return jsonify({'success': True})

@app.route('/api/execute', methods=['POST'])
def execute_command():
    data = request.json
    action = data.get('action')
    scenario_type = data.get('type', 'cpu')
    params = data.get('params', {})

    if action == 'start':
        cmd = ['python3', 'server_stress_advanced.py', scenario_type]

        if scenario_type == 'cpu':
            cmd.extend(['--cpu-percent', str(params.get('cpu_percent', 30))])
        elif scenario_type == 'mem':
            cmd.extend(['--mem-percent', str(params.get('mem_percent', 50))])
        elif scenario_type == 'disk-fill':
            cmd.extend(['--path', params.get('path', '/tmp')])
            cmd.extend(['--size', str(params.get('size', 1000))])
        elif scenario_type == 'network-delay':
            cmd.extend(['--time', str(params.get('time', 3000))])
            cmd.extend(['--interface', params.get('interface', 'eth0')])

        try:
            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                  universal_newlines=True, timeout=300)
            task_status['last_run'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            task_status['results'] = result.stdout.split('\n')
            output = result.stdout + '\n' + result.stderr if result.stderr else result.stdout
            return jsonify({'success': True, 'output': output})
        except Exception as e:
            return jsonify({'success': False, 'error': str(e)})

    elif action == 'stop':
        try:
            result = subprocess.run(['python3', 'server_stress_advanced.py', 'destroy'],
                                  stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                  universal_newlines=True, timeout=300)
            output = result.stdout + '\n' + result.stderr if result.stderr else result.stdout
            return jsonify({'success': True, 'output': output})
        except Exception as e:
            return jsonify({'success': False, 'error': str(e)})

    elif action == 'status':
        try:
            result = subprocess.run(['python3', 'server_stress_advanced.py', 'status'],
                                  stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                  universal_newlines=True, timeout=300)
            output = result.stdout + '\n' + result.stderr if result.stderr else result.stdout
            return jsonify({'success': True, 'output': output})
        except Exception as e:
            return jsonify({'success': False, 'error': str(e)})

    elif action == 'combined':
        cpu_percent = params.get('cpu_percent', 30)
        mem_percent = params.get('mem_percent', 40)

        try:
            result1 = subprocess.run(['python3', 'server_stress_advanced.py', 'cpu', '--cpu-percent', str(cpu_percent)],
                                   stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                   universal_newlines=True, timeout=300)

            result2 = subprocess.run(['python3', 'server_stress_advanced.py', 'mem', '--mem-percent', str(mem_percent)],
                                   stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                   universal_newlines=True, timeout=300)

            output = "=== CPU压测 ===\n" + result1.stdout + "\n\n=== 内存压测 ===\n" + result2.stdout
            return jsonify({'success': True, 'output': output})
        except Exception as e:
            return jsonify({'success': False, 'error': str(e)})

    elif action == 'cron_add':
        try:
            result = subprocess.run(['python3', 'cron_manager.py', 'add'],
                                  stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                  universal_newlines=True, timeout=30)
            output = result.stdout + '\n' + result.stderr if result.stderr else result.stdout
            return jsonify({'success': True, 'output': output})
        except Exception as e:
            return jsonify({'success': False, 'error': str(e)})

    elif action == 'cron_show':
        try:
            result = subprocess.run(['python3', 'cron_manager.py', 'show'],
                                  stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                  universal_newlines=True, timeout=30)
            output = result.stdout + '\n' + result.stderr if result.stderr else result.stdout
            return jsonify({'success': True, 'output': output})
        except Exception as e:
            return jsonify({'success': False, 'error': str(e)})

@app.route('/api/status')
def get_status():
    return jsonify(task_status)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)
